{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "fef8ff00-b25d-48e1-bb60-87dfa8fd0b30",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-03-13 18:56:21,301 - INFO - 从 党参 加载到 23 个有效成分\n",
      "2025-03-13 18:56:21,310 - INFO - 从 茯苓 加载到 6 个有效成分\n",
      "2025-03-13 18:56:21,320 - INFO - 从 当归 加载到 3 个有效成分\n",
      "2025-03-13 18:56:21,339 - INFO - 从 黄芪 加载到 20 个有效成分\n",
      "2025-03-13 18:56:21,348 - INFO - 从 白术 加载到 7 个有效成分\n",
      "2025-03-13 18:56:21,363 - INFO - 从 甘草 加载到 92 个有效成分\n"
     ]
    }
   ],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import networkx as nx\n",
    "import math\n",
    "import re\n",
    "import logging\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from sklearn.preprocessing import OneHotEncoder, MultiLabelBinarizer\n",
    "from sklearn.ensemble import RandomForestRegressor\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import mean_absolute_error, r2_score\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "from scipy.sparse import hstack\n",
    "from itertools import chain\n",
    "from tqdm import tqdm\n",
    "import os\n",
    "import matplotlib.font_manager as fm\n",
    "import matplotlib as mpl\n",
    "\n",
    "mpl.rcParams['font.sans-serif'] = ['SimHei']\n",
    "mpl.rcParams['axes.unicode_minus'] = False\n",
    "sns.set(style=\"whitegrid\", font_scale=1.2)\n",
    "\n",
    "try:\n",
    "    chinese_font = fm.FontProperties(fname=\"C:/Windows/Fonts/simhei.ttf\")\n",
    "    plt.rcParams[\"font.family\"] = chinese_font.get_name()\n",
    "except Exception as e:\n",
    "    logging.warning(f\"中文字体加载失败: {str(e)}\")\n",
    "\n",
    "# 配置日志\n",
    "logging.basicConfig(\n",
    "    level=logging.INFO,\n",
    "    format=\"%(asctime)s - %(levelname)s - %(message)s\",\n",
    "    handlers=[logging.FileHandler(\"analysis.log\"), logging.StreamHandler()]\n",
    ")\n",
    "\n",
    "class GeneStandardizer:\n",
    "    @staticmethod\n",
    "    def standardize(name):\n",
    "        name = re.sub(r\"\\b(mRNA|precursor|protein|subunit|chain|of)\\b\", \"\", name, flags=re.IGNORECASE)\n",
    "        name = re.sub(r\"[^a-zA-Z0-9]+\", \"_\", name)\n",
    "        return name.strip(\"_\").upper()\n",
    "\n",
    "class DataProcessor:\n",
    "    def __init__(self, ob_threshold=25, dl_threshold=0.15):\n",
    "        self.ob_thresh = ob_threshold\n",
    "        self.dl_thresh = dl_threshold\n",
    "\n",
    "    def load_components(self, file_path, herb_name):\n",
    "        try:\n",
    "            if not os.path.exists(file_path):\n",
    "                logging.error(f\"文件不存在: {file_path}\")\n",
    "                return pd.DataFrame()\n",
    "            \n",
    "            df = pd.read_csv(file_path, dtype={\"Targets\": str}, na_values=[\"N/A\", \"NA\", \"nan\"])\n",
    "            df.columns = df.columns.str.strip().str.upper()\n",
    "            df = df.rename(columns={\n",
    "                \"MOL ID\": \"mol_id\",\n",
    "                \"TARGETS\": \"targets\",\n",
    "                \"OB\": \"OB\",\n",
    "                \"DL\": \"DL\"\n",
    "            })\n",
    "            \n",
    "            df = df[\n",
    "                (pd.to_numeric(df[\"OB\"], errors=\"coerce\") >= self.ob_thresh) &\n",
    "                (pd.to_numeric(df[\"DL\"], errors=\"coerce\") >= self.dl_thresh) &\n",
    "                (~df[\"targets\"].isna())\n",
    "            ].copy()\n",
    "            \n",
    "            df[\"targets\"] = df[\"targets\"].str.split(r\",\\s*\")\n",
    "            df = df.explode(\"targets\")\n",
    "            df[\"targets\"] = df[\"targets\"].apply(\n",
    "                lambda x: GeneStandardizer.standardize(x) if pd.notna(x) else \"\"\n",
    "            )\n",
    "            df = df[df[\"targets\"] != \"\"]\n",
    "            \n",
    "            if not df.empty:\n",
    "                agg_df = df.groupby([\"mol_id\", \"OB\", \"DL\"])[\"targets\"].agg(list).reset_index()\n",
    "                agg_df[\"herb\"] = herb_name\n",
    "                logging.info(f\"从 {herb_name} 加载到 {len(agg_df)} 个有效成分\")\n",
    "                return agg_df\n",
    "            return pd.DataFrame()\n",
    "        except Exception as e:\n",
    "            logging.error(f\"文件加载失败: {file_path}\\n错误详情: {str(e)}\")\n",
    "            return pd.DataFrame()\n",
    "\n",
    "    def load_disease(self, disease_name, path):\n",
    "        try:\n",
    "            if not os.path.exists(path):\n",
    "                logging.error(f\"疾病文件不存在: {path}\")\n",
    "                return {}\n",
    "            df = pd.read_csv(path)\n",
    "            if \"gene\" not in df.columns:\n",
    "                logging.error(f\"文件 {path} 缺少 'gene' 列\")\n",
    "                return {}\n",
    "            df[\"gene\"] = df[\"gene\"].fillna(\"\").astype(str)\n",
    "            genes = df[\"gene\"].apply(GeneStandardizer.standardize).unique()\n",
    "            valid_genes = [g for g in genes if g]\n",
    "            return {disease_name: set(valid_genes)} if valid_genes else {}\n",
    "        except Exception as e:\n",
    "            logging.error(f\"疾病文件加载失败: {path}\\n错误详情: {str(e)}\")\n",
    "            return {}\n",
    "\n",
    "class AssociationAnalyzer:\n",
    "    @staticmethod\n",
    "    def calculate_association(row, disease_genes):\n",
    "        try:\n",
    "            component_genes = set(row[\"targets\"])\n",
    "            overlap = component_genes & disease_genes\n",
    "            overlap_count = len(overlap)\n",
    "            score = (row[\"OB\"]/100) * row[\"DL\"] * math.log(overlap_count + 1)\n",
    "            return pd.Series([overlap_count, \", \".join(overlap), round(score, 4)])\n",
    "        except Exception as e:\n",
    "            logging.debug(f\"计算失败: {row['mol_id']} - {str(e)}\")\n",
    "            return pd.Series([0, \"\", 0.0])\n",
    "\n",
    "class NetworkBuilder:\n",
    "    def __init__(self):\n",
    "        self.G = nx.Graph()\n",
    "    \n",
    "    def build_network(self, association_df):\n",
    "        for _, row in association_df.iterrows():\n",
    "            # 添加成分节点并设置 node_type 和 herb 属性\n",
    "            self.G.add_node(row[\"mol_id\"], \n",
    "                           node_type=\"component\", \n",
    "                           herb=row[\"herb\"])  # 关键修复：添加 herb 属性\n",
    "            # 确保疾病节点存在\n",
    "            self.G.add_node(row[\"disease\"], node_type=\"disease\")\n",
    "            # 添加成分-疾病边\n",
    "            self.G.add_edge(row[\"mol_id\"], row[\"disease\"], \n",
    "                          weight=row[\"association_score\"], \n",
    "                          relation_type=\"targets_disease\")\n",
    "            # 基因节点\n",
    "            for gene in row[\"targets\"]:\n",
    "                self.G.add_node(gene, node_type=\"gene\")\n",
    "                self.G.add_edge(row[\"mol_id\"], gene, \n",
    "                              weight=math.log(row[\"overlap_count\"] + 1),\n",
    "                              relation_type=\"has_target\")\n",
    "        return self.G\n",
    "    \n",
    "    def get_key_genes(self, top_n=10):\n",
    "        gene_degrees = [(n, d) for n, d in self.G.degree() if self.G.nodes[n].get(\"node_type\") == \"gene\"]\n",
    "        sorted_genes = sorted(gene_degrees, key=lambda x: x[1], reverse=True)[:top_n]\n",
    "        \n",
    "        results = []\n",
    "        for gene, degree in sorted_genes:\n",
    "            related_diseases = set()\n",
    "            # 遍历基因的邻居（成分节点）\n",
    "            for component in self.G.neighbors(gene):\n",
    "                if self.G.nodes[component].get(\"node_type\") == \"component\":\n",
    "                    # 遍历成分节点的邻居（疾病节点）\n",
    "                    for disease in self.G.neighbors(component):\n",
    "                        if self.G.nodes[disease].get(\"node_type\") == \"disease\":\n",
    "                            related_diseases.add(disease)\n",
    "            \n",
    "            # 修正：过滤 None 值并确保为字符串\n",
    "            related_herbs = {self.G.nodes[n].get(\"herb\", \"\") for n in self.G.neighbors(gene) \n",
    "                            if self.G.nodes[n].get(\"node_type\") == \"component\"}\n",
    "            related_herbs = [h for h in related_herbs if h]  # 移除空字符串\n",
    "            \n",
    "            results.append({\n",
    "                \"gene\": gene,\n",
    "                \"degree\": degree,\n",
    "                \"related_herbs\": \", \".join(related_herbs) if related_herbs else \"无关联中药\",\n",
    "                \"related_diseases\": \", \".join(related_diseases) if related_diseases else \"无关联疾病\"\n",
    "            })\n",
    "        return pd.DataFrame(results)\n",
    "class EnhancedVisualizer:\n",
    "    @staticmethod\n",
    "    def plot_feature_matrix(feature_matrix, save_path):\n",
    "        plt.figure(figsize=(12, 8))\n",
    "        plt.spy(feature_matrix[:50], markersize=1)\n",
    "        plt.title(\"特征矩阵结构示例（前50个成分）\", fontproperties=chinese_font)  # 指定字体\n",
    "        plt.legend(handles=[], labels=[], prop=chinese_font)  # 显式添加空图例以触发字体加载\n",
    "        plt.savefig(\"output/feature_matrix.png\", dpi=300, bbox_inches=\"tight\")\n",
    "        plt.close()    \n",
    "        \n",
    "    @staticmethod\n",
    "    def plot_heatmap(df, save_path):\n",
    "        plt.figure(figsize=(15, 20))\n",
    "        pivot = df.pivot_table(index=\"herb\", columns=\"disease\", values=\"association_score\", aggfunc=\"max\")\n",
    "        sns.heatmap(pivot, cmap=\"YlGnBu\", annot=True, fmt=\".3f\")\n",
    "        plt.title(\"中药-疾病关联热力图\", fontproperties=chinese_font)  # 显式指定字体\n",
    "        plt.savefig(save_path, dpi=300, bbox_inches=\"tight\")\n",
    "        plt.close()\n",
    "\n",
    "    @staticmethod\n",
    "    def plot_network(G, save_path):\n",
    "        plt.figure(figsize=(25, 20))\n",
    "        pos = nx.spring_layout(G, k=0.2, iterations=100)\n",
    "        node_colors = []\n",
    "        for node in G.nodes:\n",
    "            ntype = G.nodes[node].get(\"node_type\", \"other\")\n",
    "            color_map = {\n",
    "                \"disease\": \"#FF6B6B\",\n",
    "                \"herb\": \"#4CAF50\",\n",
    "                \"gene\": \"#64B5F6\",\n",
    "                \"component\": \"#CCCCCC\"\n",
    "            }\n",
    "            node_colors.append(color_map.get(ntype, \"#FFFFFF\"))\n",
    "        \n",
    "        # 绘制节点和边\n",
    "        nx.draw_networkx_nodes(G, pos, node_size=800, node_color=node_colors, alpha=0.9)\n",
    "        nx.draw_networkx_edges(G, pos, width=[G[u][v]['weight']*0.5 for u,v in G.edges], \n",
    "                             edge_color=\"#666666\", alpha=0.5)\n",
    "        \n",
    "        # 添加图例（关键修改）\n",
    "        legend_elements = [\n",
    "            plt.Line2D([0], [0], marker='o', color='w', label='疾病', \n",
    "                      markerfacecolor='#FF6B6B', markersize=15),\n",
    "            plt.Line2D([0], [0], marker='o', color='w', label='中药', \n",
    "                      markerfacecolor='#4CAF50', markersize=15),\n",
    "            plt.Line2D([0], [0], marker='o', color='w', label='基因', \n",
    "                      markerfacecolor='#64B5F6', markersize=15)\n",
    "        ]\n",
    "        plt.legend(handles=legend_elements, prop=chinese_font, loc='upper right')\n",
    "        \n",
    "        # 去除边框\n",
    "        plt.axis(\"off\")\n",
    "        plt.savefig(save_path, dpi=300, bbox_inches=\"tight\", facecolor=\"white\")  # 设置背景为白色\n",
    "        plt.close()\n",
    "\n",
    "class EnhancedAnalyzer:\n",
    "    def generate_stats(self, components, final_df, G):\n",
    "        self.stats_summary = {\n",
    "            'herb_count': components['herb'].nunique(),\n",
    "            'component_count': len(components),\n",
    "            'core_genes': len([n for n in G.nodes if G.nodes[n].get('node_type') == 'gene' \n",
    "                              and len([nbr for nbr in G.neighbors(n) if G.nodes[nbr].get('node_type') == 'disease']) >=2]),\n",
    "            'top_herb': (final_df.loc[final_df['association_score'].idxmax(), 'herb'], \n",
    "                        final_df['association_score'].max()),\n",
    "            'top_gene': max([(n, d) for n, d in G.degree() if G.nodes[n].get('node_type') == 'gene'], \n",
    "                           key=lambda x: x[1])\n",
    "        }\n",
    "    \n",
    "    def save_stats_report(self):\n",
    "        report = f\"\"\"关键统计摘要：\n",
    "1. 共分析{self.stats_summary['herb_count']}种中药的{self.stats_summary['component_count']}种成分\n",
    "2. 发现{self.stats_summary['core_genes']}个核心调控基因（与≥2种疾病相关）\n",
    "3. 关联评分最高的中药是{self.stats_summary['top_herb'][0]}，评分为{self.stats_summary['top_herb'][1]:.4f}\n",
    "4. 基因{self.stats_summary['top_gene'][0]}的网络中心性最高，数值为{self.stats_summary['top_gene'][1]}\"\"\"\n",
    "        with open(\"output/statistical_report.txt\", \"w\") as f:\n",
    "            f.write(report)\n",
    "    \n",
    "    def run_rf_analysis(self, final_df):\n",
    "        X = final_df[['OB', 'DL', 'association_score']]\n",
    "        y = final_df['overlap_count']\n",
    "        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)\n",
    "        rf = RandomForestRegressor(n_estimators=100)\n",
    "        rf.fit(X_train, y_train)\n",
    "        y_pred = rf.predict(X_test)\n",
    "        report = f\"\"\"随机森林预测结果：\n",
    "- R^2分数: {r2_score(y_test, y_pred):.4f}\n",
    "- MAE误差: {mean_absolute_error(y_test, y_pred):.4f}\n",
    "特征重要性排序：\n",
    "{pd.Series(rf.feature_importances_, index=X.columns).sort_values(ascending=False).to_string()}\"\"\"\n",
    "        with open(\"output/ml_report.txt\", \"w\") as f:\n",
    "            f.write(report)\n",
    "\n",
    "# 配置路径\n",
    "COMPONENT_FILES = {\n",
    "    \"党参\": \"Components_Corrected党参.csv\",\n",
    "    \"茯苓\": \"Components_Corrected茯苓.csv\",\n",
    "    \"当归\": \"Components_Corrected当归.csv\",\n",
    "    \"黄芪\": \"Components_Corrected黄芪.csv\",\n",
    "    \"白术\": \"Components_Corrected白术.csv\",\n",
    "    \"甘草\": \"Components_Corrected甘草.csv\"\n",
    "}\n",
    "\n",
    "DISEASE_FILES = {\n",
    "    \"Asthma\": \"Asthma.csv\",\n",
    "    \"Hepatitis_B\": \"Hepatitis_B.csv\",\n",
    "    \"Adult_Onset_Diabetes_Mellitus\": \"Adult_Onset_Diabetes_Mellitus.csv\"\n",
    "}\n",
    "\n",
    "def main():\n",
    "    os.makedirs(\"output\", exist_ok=True)\n",
    "    processor = DataProcessor()\n",
    "    analyzer = EnhancedAnalyzer()\n",
    "    network = NetworkBuilder()\n",
    "    \n",
    "    try:\n",
    "        # 加载数据\n",
    "        components = []\n",
    "        for herb, path in COMPONENT_FILES.items():\n",
    "            df = processor.load_components(path, herb)\n",
    "            if not df.empty:\n",
    "                components.append(df)\n",
    "        component_df = pd.concat(components).reset_index(drop=True)\n",
    "        \n",
    "        diseases = {}\n",
    "        for disease_name, path in DISEASE_FILES.items():\n",
    "            data = processor.load_disease(disease_name, path)\n",
    "            if data:\n",
    "                diseases.update(data)\n",
    "        \n",
    "        # 关联度计算\n",
    "        association_results = []\n",
    "        for disease_name, disease_genes in diseases.items():\n",
    "            temp_df = component_df.copy()\n",
    "            temp_df[[\"overlap_count\", \"overlap_genes\", \"association_score\"]] = temp_df.apply(\n",
    "                lambda x: AssociationAnalyzer.calculate_association(x, disease_genes), axis=1\n",
    "            )\n",
    "            temp_df[\"disease\"] = disease_name\n",
    "            filtered_df = temp_df[temp_df[\"association_score\"] > 0]\n",
    "            if not filtered_df.empty:\n",
    "                association_results.append(filtered_df)\n",
    "        \n",
    "        final_df = pd.concat(association_results)\n",
    "        final_df.to_csv(\"output/component_disease_association.csv\", index=False)\n",
    "        \n",
    "        # 构建网络\n",
    "        G = network.build_network(final_df)\n",
    "        nx.write_gexf(G, \"output/tcm_network.gexf\")\n",
    "        key_genes_df = network.get_key_genes(top_n=10)\n",
    "        key_genes_df.to_csv(\"output/core_regulatory_genes.csv\", index=False)\n",
    "        \n",
    "        # 可视化\n",
    "        EnhancedVisualizer.plot_heatmap(final_df, \"output/herb_disease_heatmap.png\")\n",
    "        EnhancedVisualizer.plot_network(G, \"output/enhanced_network.png\")\n",
    "        \n",
    "        # 统计与机器学习\n",
    "        analyzer.generate_stats(component_df, final_df, G)\n",
    "        analyzer.save_stats_report()\n",
    "        analyzer.run_rf_analysis(final_df)\n",
    "        \n",
    "    except Exception as e:\n",
    "        logging.error(f\"流程异常终止: {str(e)}\")\n",
    "        raise\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b27a5ec8-e114-4ef6-acfa-0c005449f003",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
